#include "../include/http_server.h"
#include "../include/faiss_index.h"
#include "../include/hnswlib_index.h"
#include "../include/index_factory.h"
#include "../include/logger.h"
#include "../include/constants.h"
#include "../include/rapidjson/document.h"
#include "../include/rapidjson/writer.h"
#include "../include/rapidjson/stringbuffer.h"
#include <iostream>

HttpServer::HttpServer(const std::string& host, int port) : host(host), port(port) {
    server.Post("/search", [this](const httplib::Request& req, httplib::Response& res) {
      searchHandler(req, res);
     });

    server.Post("/insert", [this](const httplib::Request& req, httplib::Response& res){
      insertHandler(req, res);
    });
}

void HttpServer::start() {
    server.listen(host.c_str(), port);
}

bool HttpServer::isRequestValid(const rapidjson::Document& json_request, CheckType check_type) {
  switch (check_type) {
      case CheckType::SEARCH:
          return json_request.HasMember(REQUEST_VECTORS) && 
                 json_request.HasMember(REQUEST_K) &&
                 (!json_request.HasMember(REQUEST_INDEX_TYPE) || json_request[REQUEST_INDEX_TYPE].IsString());
      case CheckType::INSERT:
          return json_request.HasMember(REQUEST_VECTORS) && 
                 json_request.HasMember(REQUEST_ID) &&
                 (!json_request.HasMember(REQUEST_INDEX_TYPE) || json_request[REQUEST_INDEX_TYPE].IsString());
      default:
          return false;
  }
}

IndexFactory::IndexType HttpServer::getIndexTypeFromRequest(const rapidjson::Document& json_request) {
    if (json_request.HasMember(REQUEST_INDEX_TYPE)) {
        std::string index_type_str = json_request[REQUEST_INDEX_TYPE].GetString();
        if (index_type_str == INDEX_TYPE_FLAT) {
            return IndexFactory::IndexType::FLAT;
        } else if (index_type_str == INDEX_TYPE_HNSW) {
            return IndexFactory::IndexType::HNSW;
        }
    }
    return IndexFactory::IndexType::UNKOWN;
}

void HttpServer::searchHandler(const httplib::Request& req, httplib::Response& res) {
  GlobalLogger->debug("Received search request");

  // parse json request
  rapidjson::Document json_request;
  json_request.Parse(req.body.c_str());

  // print user's input params
  GlobalLogger->info("Search request parameters: {}", req.body);

  // check JSON document is valid object
  if (!json_request.IsObject()) {
      GlobalLogger->error("Invalid JSON request");
      res.status = 400;
      setErrorJsonResponse(res, RESPONCE_RETCODE_ERROR, "Invalid JSON request");
      return ;
  }

  // check the request is valid
  if (!isRequestValid(json_request, CheckType::SEARCH)) {
      GlobalLogger->error("Missing vectors or k parameter in the request");
      res.status = 400;
      setErrorJsonResponse(res, RESPONCE_RETCODE_ERROR, "Missing vectors or k parameter in the request");
      return ;
  }

  // get the search params
  std::vector<float> query;
  for (const auto& q : json_request[REQUEST_VECTORS].GetArray()) {
      query.emplace_back(q.GetFloat());
  }

  int k = json_request[REQUEST_K].GetInt();

  GlobalLogger->debug("Query parameters: k = {}", k);

  // get the index type
  IndexFactory::IndexType indexType = getIndexTypeFromRequest(json_request);

  // if indexType is UNKOWN, return 400 error
  if (indexType == IndexFactory::IndexType::UNKOWN) {
      GlobalLogger->error("Invalid indexType parameter in the request");
      res.status = 400;
      setErrorJsonResponse(res, RESPONCE_RETCODE_ERROR, "Invalid indexType parameter in the request");
      return ;
  }

  void* index = getGlobalIndexFactory()->getIndex(indexType);

  // init index object according the indexType and revoke search_vector
  std::pair<std::vector<long>, std::vector<float>> results;
  switch(indexType) {
      case IndexFactory::IndexType::FLAT: {
          FaissIndex* faissIndex = static_cast<FaissIndex*>(index);
          results = faissIndex->search_vectors(query, k);
          break;
      }
      case IndexFactory::IndexType::HNSW: {
          HNSWLibIndex* hnswIndex = static_cast<HNSWLibIndex*>(index); 
          results = hnswIndex->search_vectors(query, k);
          break;
      }
      default:
        break;
      
  }

  // trans results to json
  rapidjson::Document json_responce;
  json_responce.SetObject();
  rapidjson::Document::AllocatorType& allocator = json_responce.GetAllocator();

  // check if the search results are valid
  bool valid_results = false;
  rapidjson::Value vectors(rapidjson::kArrayType);
  rapidjson::Value distances(rapidjson::kArrayType);
  for (size_t i = 0; i < results.first.size(); ++i) {
      if (results.first[i] != -1) {
          valid_results = true;
          vectors.PushBack(results.first[i], allocator);
          distances.PushBack(results.second[i], allocator);
      }
  }

  if (valid_results) {
      json_responce.AddMember(RESPONCE_VECTORS, vectors, allocator);
      json_responce.AddMember(RESPONCE_DISTANCES, distances, allocator);
  }

  // set Response
  json_responce.AddMember(RESPONCE_RETCODE, RESPONCE_RETCODE_SUCCESS, allocator); 
  setJsonResponse(json_responce, res);
}

void HttpServer::insertHandler(const httplib::Request& req, httplib::Response& res) {
  GlobalLogger->debug("Received insert request");

  // parse json request
  rapidjson::Document json_request;
  json_request.Parse(req.body.c_str());

  // print user's input params
  GlobalLogger->info("Insert request parameters: {}", req.body);

  // check JSON document is valid object
  if (!json_request.IsObject()) {
      GlobalLogger->error("Invalid JSON request");
      res.status = 400;
      setErrorJsonResponse(res, RESPONCE_RETCODE_ERROR, "Invalid JSON request");
      return ;
  }

  // check the request is valid
  if (!isRequestValid(json_request, CheckType::INSERT)) {
      GlobalLogger->error("Missing vectors or id parameter in the request");
      res.status = 400;
      setErrorJsonResponse(res, RESPONCE_RETCODE_ERROR, "Missing vectors or id parameter in the request");
      return ;
  }

  // get the insert params
  std::vector<float> data;
  for (const auto& d : json_request[REQUEST_VECTORS].GetArray()) {
      data.emplace_back(d.GetFloat());
  }

  uint64_t label = json_request[REQUEST_ID].GetUint64();

  GlobalLogger->debug("Query parameters: label = {}", label);

  // get the index type
  IndexFactory::IndexType indexType = getIndexTypeFromRequest(json_request);

  // if indexType is UNKOWN, return 400 error
  if (indexType == IndexFactory::IndexType::UNKOWN) {
      GlobalLogger->error("Invalid indexType parameter in the request");
      res.status = 400;
      setErrorJsonResponse(res, RESPONCE_RETCODE_ERROR, "Invalid indexType parameter in the request");
      return ;
  }

  void* index = getGlobalIndexFactory()->getIndex(indexType);

  // init index object according the indexType and revoke insert_vectors
  switch(indexType) {
      case IndexFactory::IndexType::FLAT: {
          FaissIndex* faissIndex = static_cast<FaissIndex*>(index);
          faissIndex->insert_vectors(data, label);
          break;
      }
      case IndexFactory::IndexType::HNSW: {
          HNSWLibIndex* hnswIndex = static_cast<HNSWLibIndex*>(index); 
          hnswIndex->insert_vectors(data, label);
          break;
      } 
      default:
        break;
      
  }

  // set Response
  rapidjson::Document json_responce;
  json_responce.SetObject();
  rapidjson::Document::AllocatorType& allocator = json_responce.GetAllocator();
  json_responce.AddMember(RESPONCE_RETCODE, RESPONCE_RETCODE_SUCCESS, allocator); 
  setJsonResponse(json_responce, res);
}

void HttpServer::setJsonResponse(const rapidjson::Document& json_responce, httplib::Response& res) {
  rapidjson::StringBuffer buffer;
  rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
  json_responce.Accept(writer);
  res.set_content(buffer.GetString(), RESPONCE_CONTENT_TYPE_JSON);
}

void HttpServer::setErrorJsonResponse(httplib::Response& res, int error_code, const std::string& errorMsg) {
  rapidjson::Document json_responce;
  json_responce.SetObject();
  rapidjson::Document::AllocatorType& allocator = json_responce.GetAllocator();
  json_responce.AddMember(RESPONCE_RETCODE, error_code, allocator);
  json_responce.AddMember(RESPONCE_ERROR_MSG, rapidjson::StringRef(errorMsg.c_str()), allocator);
  setJsonResponse(json_responce, res);
}
